package state

import (
	"context"
	"fmt"

	"github.com/mcp-ecosystem/mcp-gateway/internal/common/cnst"
	"github.com/mcp-ecosystem/mcp-gateway/internal/common/config"
	"github.com/mcp-ecosystem/mcp-gateway/internal/core/mcpproxy"
	"github.com/mcp-ecosystem/mcp-gateway/internal/template"
	"github.com/mcp-ecosystem/mcp-gateway/pkg/mcp"
	"go.uber.org/zap"
)

type (
	uriPrefix string
	toolName  string

	// State contains all the read-only shared state
	State struct {
		rawConfigs []*config.MCPConfig
		runtime    map[uriPrefix]runtimeUnit
		metrics    metrics
	}

	runtimeUnit struct {
		protoType cnst.ProtoType
		router    *config.RouterConfig
		server    *config.ServerConfig
		mcpServer *config.MCPServerConfig
		transport mcpproxy.Transport

		tools       map[toolName]*config.ToolConfig
		toolSchemas []mcp.ToolSchema
	}

	metrics struct {
		totalTools      int
		missingTools    int
		httpServers     int
		mcpServers      int
		idleHTTPServers int
		idleMCPServers  int
	}
)

func NewState() *State {
	return &State{
		rawConfigs: make([]*config.MCPConfig, 0),
		runtime:    make(map[uriPrefix]runtimeUnit),
		metrics:    metrics{},
	}
}

// BuildStateFromConfig creates a new State from the given configuration
func BuildStateFromConfig(ctx context.Context, cfgs []*config.MCPConfig, oldState *State, logger *zap.Logger) (*State, error) {
	// Create new state
	newState := NewState()
	newState.rawConfigs = cfgs

	for _, cfg := range cfgs {
		toolMap := make(map[toolName]*config.ToolConfig)
		// Initialize tool map and list for MCP servers
		for _, tool := range cfg.Tools {
			newState.metrics.totalTools++
			toolMap[toolName(tool.Name)] = &tool
		}

		// Build prefix to tools mapping for MCP servers
		prefixMap := make(map[string]string)
		for _, router := range cfg.Routers {
			prefixMap[router.Server] = router.Prefix
			newState.setRouter(router.Prefix, &router)
		}

		// Process regular HTTP servers
		for _, server := range cfg.Servers {
			newState.metrics.httpServers++
			prefix, ok := prefixMap[server.Name]
			if !ok {
				newState.metrics.idleHTTPServers++
				logger.Warn("failed to find prefix for server", zap.String("server", server.Name))
				continue
			}

			var (
				allowedToolSchemas []mcp.ToolSchema
				allowedTools       = make(map[toolName]*config.ToolConfig)
			)
			for _, ss := range server.AllowedTools {
				tool, ok := toolMap[toolName(ss)]
				if ok {
					allowedToolSchemas = append(allowedToolSchemas, tool.ToToolSchema())
					allowedTools[toolName(ss)] = tool
				} else {
					newState.metrics.missingTools++
					logger.Warn("failed to find allowed tool for server", zap.String("server", server.Name),
						zap.String("tool", ss))
				}
			}
			runtime := newState.getRuntime(prefix)
			runtime.protoType = cnst.BackendProtoHttp
			runtime.server = &server
			runtime.tools = allowedTools
			runtime.toolSchemas = allowedToolSchemas
			newState.runtime[uriPrefix(prefix)] = runtime
		}

		// Process MCP servers
		for _, mcpServer := range cfg.McpServers {
			newState.metrics.mcpServers++
			prefix, exists := prefixMap[mcpServer.Name]
			if !exists {
				newState.metrics.idleMCPServers++
				logger.Warn("failed to find prefix for mcp server", zap.String("server", mcpServer.Name))
				continue // Skip MCP servers without router prefix
			}

			runtime := newState.getRuntime(prefix)
			runtime.mcpServer = &mcpServer

			// Check if we already have transport with the same configuration
			var transport mcpproxy.Transport
			if oldState != nil {
				if oldRuntime, exists := oldState.runtime[uriPrefix(prefix)]; exists {
					// Compare configurations to see if we need to create a new transport
					oldConfig := oldRuntime.mcpServer
					if oldConfig != nil && oldConfig.Type == mcpServer.Type &&
						oldConfig.Command == mcpServer.Command &&
						oldConfig.URL == mcpServer.URL &&
						len(oldConfig.Args) == len(mcpServer.Args) {
						// Compare args
						argsMatch := true
						for i, arg := range oldConfig.Args {
							if arg != mcpServer.Args[i] {
								argsMatch = false
								break
							}
						}
						if argsMatch {
							// Reuse existing transport
							transport = oldRuntime.transport
						}
					}
				}
			}

			// Create new transport if needed
			if transport == nil {
				var err error
				transport, err = mcpproxy.NewTransport(mcpServer)
				if err != nil {
					return nil, fmt.Errorf("failed to create transport for server %s: %w", mcpServer.Name, err)
				}
			}

			// Handle server startup based on policy and preinstalled flag
			if mcpServer.Policy == cnst.PolicyOnStart {
				// If PolicyOnStart is set, just start the server and keep it running
				go startMCPServer(ctx, logger, prefix, mcpServer, transport, false)
			} else if mcpServer.Preinstalled {
				// If Preinstalled is set but not PolicyOnStart, verify installation by starting and stopping
				go startMCPServer(ctx, logger, prefix, mcpServer, transport, true)
			}
			runtime.transport = transport

			// Map protocol type based on server type
			switch mcpServer.Type {
			case cnst.BackendProtoStdio.String():
				runtime.protoType = cnst.BackendProtoStdio
			case cnst.BackendProtoSSE.String():
				runtime.protoType = cnst.BackendProtoSSE
			case cnst.BackendProtoStreamable.String():
				runtime.protoType = cnst.BackendProtoStreamable
			}
			newState.runtime[uriPrefix(prefix)] = runtime
		}
	}

	if oldState != nil {
		for prefix, oldRuntime := range oldState.runtime {
			if _, stillExists := newState.runtime[prefix]; !stillExists {
				if oldRuntime.mcpServer == nil {
					continue
				}
				if oldRuntime.transport == nil {
					logger.Info("transport already stopped", zap.String("prefix", string(prefix)),
						zap.String("command", oldRuntime.mcpServer.Command), zap.Strings("args", oldRuntime.mcpServer.Args))
					continue
				}
				logger.Info("shutting down unused transport", zap.String("prefix", string(prefix)),
					zap.String("command", oldRuntime.mcpServer.Command), zap.Strings("args", oldRuntime.mcpServer.Args))
				if err := oldRuntime.transport.Stop(ctx); err != nil {
					logger.Warn("failed to close old transport", zap.String("prefix", string(prefix)),
						zap.Error(err), zap.String("command", oldRuntime.mcpServer.Command),
						zap.Strings("args", oldRuntime.mcpServer.Args))
				}
			}
		}
	}

	return newState, nil
}

func startMCPServer(ctx context.Context, logger *zap.Logger, prefix string, mcpServer config.MCPServerConfig,
	transport mcpproxy.Transport, needStop bool) {
	// If Preinstalled is set but not PolicyOnStart, verify installation by starting and stopping
	if transport.IsRunning() {
		logger.Info("server already started, don't need to preinstall",
			zap.String("prefix", prefix),
			zap.String("command", mcpServer.Command),
			zap.Strings("args", mcpServer.Args))
		return
	}

	if err := transport.Start(ctx, template.NewContext()); err != nil {
		logger.Error("failed to start server for preinstall",
			zap.String("prefix", prefix),
			zap.String("command", mcpServer.Command),
			zap.Strings("args", mcpServer.Args),
			zap.Error(err))
	} else {
		logger.Info("server started for preinstall",
			zap.String("prefix", prefix),
			zap.String("command", mcpServer.Command),
			zap.Strings("args", mcpServer.Args))

		if needStop {
			// Stop the server after successful start
			if err := transport.Stop(ctx); err != nil {
				logger.Error("failed to stop server for preinstall",
					zap.String("prefix", prefix),
					zap.String("command", mcpServer.Command),
					zap.Strings("args", mcpServer.Args),
					zap.Error(err))
			} else {
				logger.Info("server stopped for preinstall",
					zap.String("prefix", prefix),
					zap.String("command", mcpServer.Command),
					zap.Strings("args", mcpServer.Args))
			}
		}
	}
}
